#pragma once

#include "updater_common.h"

namespace mdtx::online::updater
{

    // TODO: enable Kahan compensation for arma vectors

    struct smav_updater
    {
        size_t n;
        arma::rowvec m;

        explicit smav_updater(int nvar)
            : n(0),
              m(nvar, arma::fill::zeros)
        {
        }

        void insert(const arma::rowvec &x)
        {
            n += 1;
            m += (x - m) / n;
        }

        void remove(const arma::rowvec &x)
        {
            n -= 1;
            m -= (x - m) / n;
        }

        void roll(const arma::rowvec &x, const arma::rowvec &y)
        {
            m += (x - y) / n;
        }
    };

    struct wmav_updater
    {
        size_t n, w;
        arma::rowvec xsum, wsum;

        explicit wmav_updater(int nvar)
            : n(0),
              w(0),
              xsum(nvar, arma::fill::zeros),
              wsum(nvar, arma::fill::zeros)
        {
        }

        void insert(const arma::rowvec &x)
        {
            n += 1;
            w += n;
            wsum += x * n;
            xsum += x;
        }

        void remove(const arma::rowvec &x)
        {
            w -= n;
            n -= 1;
            wsum -= xsum;
            xsum -= x;
        }

        void roll(const arma::colvec &x, const arma::colvec &y)
        {
            wsum += x * n - xsum;
            xsum += x - y;
        }
    };

    struct emav_updater
    {
        bool init;
        double a;
        arma::rowvec m;

        emav_updater(int nvar, double alpha)
            : init(false),
              a(alpha),
              m(nvar, arma::fill::zeros)
        {
        }

        void insert(const arma::rowvec &x)
        {
            if (init)
            {
                m += a * (x - m);
            }
            else
            {
                m = x;
                init = true;
            }
        }

        void remove(const arma::rowvec &x)
        {
            x;
        }

        void roll(const arma::colvec &x, const arma::colvec &y)
        {
            insert(x);
            y;
        }
    };

}
